#!/usr/bin/env python2
from __future__ import division

from pylab import *
from numpy import * #unnecessary, just for pydev auto-complete
from matching_pursuit import *
from fJ_generator import *
from scipy import optimize
from scikits.audiolab import wavread
from scikits.audiolab import play
from levenberg_marquardt import ooopti
#from enthought.mayavi.mlab import *
import os


class rt:
    def __init__(self, block_size):
        x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10 = sympy.symbols(['x0', 'x1', 'x2', 'x3', 'x4', 'x5', 'x6', 'x7', 'x8', 'x9', 'x10'])
        t = sympy.Symbol('t')
#        self.fJ_generator = fJ_generator(sympy.log(1 + sympy.exp(100 * (x0 + x1 * t + x2 * t ** 2))) / 100 * sympy.cos((x3 + x4 * t + x5 * t ** 2)), [x0, x1, x2, x3, x4, x5 ], {t:linspace(0, 1, block_size, endpoint=False)})
        self.fJ_generator = fJ_generator(sympy.log(1 + sympy.exp(10 * (x0 + x1 * t +x2* t**2))) / 10 * sympy.cos((x3 + x4 * t + x5 * t ** 2)), [x0, x1, x2, x3, x4, x5 ], {t:linspace(0, 1, block_size, endpoint=False)})
        self.mp = MatchingPursuit(block_size, 12.25 / 4, 120 * 10, 120)

    def rt(self, data):
        ww = []
        x = array([]) 
        for i in xrange(10):
            print i, 'th iteration'
            diff = data - self.fJ_generator.f(x)
            [A, w, p] = self.mp.matching_pursuit(diff)
            ww += [w]
            (add, success) = optimize.leastsq(self.fJ_generator.f, [A, 0, 0, p, w, 0], args=(diff), Dfun=self.fJ_generator.J, full_output=0)
#            (add, success) = optimize.leastsq(self.fJ_generator.f, [A, 0, p, w, 0], args=(diff), Dfun=self.fJ_generator.J, full_output=0)
            
#            add= ooopti(self.fJ_generator.f, [A, 0, 0, p, w, 0, 0, 0], args=(diff), Dfun=self.fJ_generator.J)
            x = concatenate([x, add])
            try:
                (x, success, s2, s3, s4) = optimize.leastsq(self.fJ_generator.f, x, args=(data), Dfun=self.fJ_generator.J, full_output=1)
                
#                x = ooopti(self.fJ_generator.f, x, args=(data), Dfun=self.fJ_generator.J)
            except:
                print success,s2,s3,s4
                break

            print ww
        return x
    
    
if __name__ == '__main__':
    windowsize = 4410
    wav, fs, enc = wavread("/home/kkb110/music/g.wav")
    xxxxx = rt(windowsize)
    #data = xxxxx.fJ_generator.f([1, 0, 0, 0, 20 * pi * 2, 0]+[3,0,0,2,67*pi*2,0])
    data = wav.T[0, 0.8 * 44100:0.8 * 44100 + windowsize]
    x = xxxxx.rt(data)
    
    #p3d(x.reshape(-1, 10))
    for i in range(5):
        plot(xxxxx.fJ_generator.f(x.reshape(-1, 5)[i]))
        show()
    plot(data)
    plot(xxxxx.fJ_generator.f(x), 'r--')
    plot(data - xxxxx.fJ_generator.f(x))
    show()
    while(True):
        i = input("1 data, 2 x, 3, data-x")
        if i == 1:
            play(data)
        if i == 2:
            play(xxxxx.fJ_generator.f(x))
        if i == 3:
            play(data - xxxxx.fJ_generator.f(x))
    
    input("sdf")
